{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "NTaMBjChz7m5"
   },
   "source": [
    "# Google ADK + Opik Integration Cookbook\n",
    "\n",
    "This notebook demonstrates how to integrate Google's Agent Development Kit (ADK) with Opik for comprehensive tracing and observability. We'll cover three key integration patterns:\n",
    "\n",
    "1. **Basic Agent Example** - Simple single-agent setup with Opik tracing\n",
    "2. **Multi-Agent Example** - Complex multi-agent workflow showing hierarchical tracing\n",
    "3. **Hybrid Tracing** - Combining Opik decorators with ADK callbacks for comprehensive observability\n",
    "\n",
    "You will need:\n",
    "\n",
    "1. A Comet account, for seeing Opik visualizations (free!) - [comet.com](https://comet.com)\n",
    "2. An OpenAI account, for using gpt-4o model - [platform.openai.com/settings/organization/api-keys](https://platform.openai.com/settings/organization/api-keys)\n",
    "3. Google ADK installed and configured\n",
    "\n",
    "This example will use:\n",
    "\n",
    "- [google-adk](https://pypi.org/project/google-adk/) for agent development\n",
    "- [opik](https://pypi.org/project/opik/) for tracing and observability\n",
    "- OpenAI's gpt-4o model through LiteLLM"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yM1lU0dBBnJs"
   },
   "source": [
    "## Setup\n",
    "\n",
    "Install the required packages:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "2Tx6HwuU1rB4"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: opik in /home/mavrick/anaconda3/lib/python3.12/site-packages (1.7.42)\n",
      "Requirement already satisfied: google-adk in /home/mavrick/anaconda3/lib/python3.12/site-packages (1.5.0)\n",
      "Requirement already satisfied: litellm in /home/mavrick/anaconda3/lib/python3.12/site-packages (1.74.0)\n",
      "Requirement already satisfied: boto3-stubs>=1.34.110 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from boto3-stubs[bedrock-runtime]>=1.34.110->opik) (1.34.133)\n",
      "Requirement already satisfied: click in /home/mavrick/anaconda3/lib/python3.12/site-packages (from opik) (8.1.8)\n",
      "Requirement already satisfied: httpx in /home/mavrick/anaconda3/lib/python3.12/site-packages (from opik) (0.28.1)\n",
      "Requirement already satisfied: rapidfuzz<4.0.0,>=3.0.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from opik) (3.13.0)\n",
      "Requirement already satisfied: openai<2.0.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from opik) (1.93.0)\n",
      "Requirement already satisfied: pydantic-settings!=2.9.0,<3.0.0,>=2.0.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from opik) (2.6.1)\n",
      "Requirement already satisfied: pydantic<3.0.0,>=2.0.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from opik) (2.11.7)\n",
      "Requirement already satisfied: pytest in /home/mavrick/anaconda3/lib/python3.12/site-packages (from opik) (8.3.5)\n",
      "Requirement already satisfied: rich in /home/mavrick/anaconda3/lib/python3.12/site-packages (from opik) (13.9.4)\n",
      "Requirement already satisfied: sentry_sdk>=2.0.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from opik) (2.14.0)\n",
      "Requirement already satisfied: tenacity in /home/mavrick/anaconda3/lib/python3.12/site-packages (from opik) (8.5.0)\n",
      "Requirement already satisfied: tqdm in /home/mavrick/anaconda3/lib/python3.12/site-packages (from opik) (4.67.1)\n",
      "Requirement already satisfied: uuid6 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from opik) (2025.0.0)\n",
      "Requirement already satisfied: jinja2 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from opik) (3.1.6)\n",
      "Requirement already satisfied: PyYAML>=6.0.2 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-adk) (6.0.2)\n",
      "Requirement already satisfied: anyio>=4.9.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-adk) (4.9.0)\n",
      "Requirement already satisfied: authlib>=1.5.1 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-adk) (1.6.0)\n",
      "Requirement already satisfied: fastapi>=0.115.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-adk) (0.115.12)\n",
      "Requirement already satisfied: google-api-python-client>=2.157.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-adk) (2.174.0)\n",
      "Requirement already satisfied: google-cloud-aiplatform>=1.95.1 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-cloud-aiplatform[agent-engines]>=1.95.1->google-adk) (1.100.0)\n",
      "Requirement already satisfied: google-cloud-secret-manager>=2.22.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-adk) (2.24.0)\n",
      "Requirement already satisfied: google-cloud-speech>=2.30.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-adk) (2.33.0)\n",
      "Requirement already satisfied: google-cloud-storage<3.0.0,>=2.18.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-adk) (2.19.0)\n",
      "Requirement already satisfied: google-genai>=1.21.1 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-adk) (1.21.1)\n",
      "Requirement already satisfied: graphviz>=0.20.2 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-adk) (0.21)\n",
      "Requirement already satisfied: mcp>=1.8.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-adk) (1.9.4)\n",
      "Requirement already satisfied: opentelemetry-api>=1.31.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-adk) (1.34.1)\n",
      "Requirement already satisfied: opentelemetry-exporter-gcp-trace>=1.9.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-adk) (1.9.0)\n",
      "Requirement already satisfied: opentelemetry-sdk>=1.31.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-adk) (1.34.1)\n",
      "Requirement already satisfied: python-dateutil>=2.9.0.post0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-adk) (2.9.0.post0)\n",
      "Requirement already satisfied: python-dotenv>=1.0.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-adk) (1.1.1)\n",
      "Requirement already satisfied: requests>=2.32.4 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-adk) (2.32.4)\n",
      "Requirement already satisfied: sqlalchemy>=2.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-adk) (2.0.15)\n",
      "Requirement already satisfied: starlette>=0.46.2 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-adk) (0.46.2)\n",
      "Requirement already satisfied: typing-extensions<5,>=4.5 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-adk) (4.12.2)\n",
      "Requirement already satisfied: tzlocal>=5.3 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-adk) (5.3.1)\n",
      "Requirement already satisfied: uvicorn>=0.34.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-adk) (0.35.0)\n",
      "Requirement already satisfied: websockets>=15.0.1 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-adk) (15.0.1)\n",
      "Requirement already satisfied: aiohttp>=3.10 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from litellm) (3.11.16)\n",
      "Requirement already satisfied: importlib-metadata>=6.8.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from litellm) (8.5.0)\n",
      "Requirement already satisfied: jsonschema<5.0.0,>=4.22.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from litellm) (4.23.0)\n",
      "Requirement already satisfied: tiktoken>=0.7.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from litellm) (0.7.0)\n",
      "Requirement already satisfied: tokenizers in /home/mavrick/anaconda3/lib/python3.12/site-packages (from litellm) (0.21.1)\n",
      "Requirement already satisfied: aiohappyeyeballs>=2.3.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from aiohttp>=3.10->litellm) (2.4.4)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from aiohttp>=3.10->litellm) (1.2.0)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from aiohttp>=3.10->litellm) (24.3.0)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from aiohttp>=3.10->litellm) (1.5.0)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from aiohttp>=3.10->litellm) (6.1.0)\n",
      "Requirement already satisfied: propcache>=0.2.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from aiohttp>=3.10->litellm) (0.3.1)\n",
      "Requirement already satisfied: yarl<2.0,>=1.17.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from aiohttp>=3.10->litellm) (1.18.0)\n",
      "Requirement already satisfied: idna>=2.8 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from anyio>=4.9.0->google-adk) (3.7)\n",
      "Requirement already satisfied: sniffio>=1.1 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from anyio>=4.9.0->google-adk) (1.3.0)\n",
      "Requirement already satisfied: cryptography in /home/mavrick/anaconda3/lib/python3.12/site-packages (from authlib>=1.5.1->google-adk) (44.0.1)\n",
      "Requirement already satisfied: botocore-stubs in /home/mavrick/anaconda3/lib/python3.12/site-packages (from boto3-stubs>=1.34.110->boto3-stubs[bedrock-runtime]>=1.34.110->opik) (1.38.30)\n",
      "Requirement already satisfied: types-s3transfer in /home/mavrick/anaconda3/lib/python3.12/site-packages (from boto3-stubs>=1.34.110->boto3-stubs[bedrock-runtime]>=1.34.110->opik) (0.13.0)\n",
      "Requirement already satisfied: mypy-boto3-bedrock-runtime<1.35.0,>=1.34.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from boto3-stubs[bedrock-runtime]>=1.34.110->opik) (1.34.149)\n",
      "Requirement already satisfied: httplib2<1.0.0,>=0.19.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-api-python-client>=2.157.0->google-adk) (0.22.0)\n",
      "Requirement already satisfied: google-auth!=2.24.0,!=2.25.0,<3.0.0,>=1.32.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-api-python-client>=2.157.0->google-adk) (2.40.3)\n",
      "Requirement already satisfied: google-auth-httplib2<1.0.0,>=0.2.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-api-python-client>=2.157.0->google-adk) (0.2.0)\n",
      "Requirement already satisfied: google-api-core!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.0,<3.0.0,>=1.31.5 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-api-python-client>=2.157.0->google-adk) (2.25.1)\n",
      "Requirement already satisfied: uritemplate<5,>=3.0.1 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-api-python-client>=2.157.0->google-adk) (4.2.0)\n",
      "Requirement already satisfied: proto-plus<2.0.0,>=1.22.3 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-cloud-aiplatform>=1.95.1->google-cloud-aiplatform[agent-engines]>=1.95.1->google-adk) (1.26.1)\n",
      "Requirement already satisfied: protobuf!=4.21.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5,<7.0.0,>=3.20.2 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-cloud-aiplatform>=1.95.1->google-cloud-aiplatform[agent-engines]>=1.95.1->google-adk) (4.25.8)\n",
      "Requirement already satisfied: packaging>=14.3 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-cloud-aiplatform>=1.95.1->google-cloud-aiplatform[agent-engines]>=1.95.1->google-adk) (24.2)\n",
      "Requirement already satisfied: google-cloud-bigquery!=3.20.0,<4.0.0,>=1.15.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-cloud-aiplatform>=1.95.1->google-cloud-aiplatform[agent-engines]>=1.95.1->google-adk) (3.34.0)\n",
      "Requirement already satisfied: google-cloud-resource-manager<3.0.0,>=1.3.3 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-cloud-aiplatform>=1.95.1->google-cloud-aiplatform[agent-engines]>=1.95.1->google-adk) (1.14.2)\n",
      "Requirement already satisfied: shapely<3.0.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-cloud-aiplatform>=1.95.1->google-cloud-aiplatform[agent-engines]>=1.95.1->google-adk) (2.0.6)\n",
      "Requirement already satisfied: docstring_parser<1 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-cloud-aiplatform>=1.95.1->google-cloud-aiplatform[agent-engines]>=1.95.1->google-adk) (0.16)\n",
      "Requirement already satisfied: cloudpickle<4.0,>=3.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-cloud-aiplatform[agent-engines]>=1.95.1->google-adk) (3.0.0)\n",
      "Requirement already satisfied: google-cloud-trace<2 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-cloud-aiplatform[agent-engines]>=1.95.1->google-adk) (1.16.2)\n",
      "Requirement already satisfied: google-cloud-logging<4 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-cloud-aiplatform[agent-engines]>=1.95.1->google-adk) (3.12.1)\n",
      "Requirement already satisfied: grpc-google-iam-v1<1.0.0,>=0.14.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-cloud-secret-manager>=2.22.0->google-adk) (0.14.2)\n",
      "Requirement already satisfied: google-cloud-core<3.0dev,>=2.3.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-cloud-storage<3.0.0,>=2.18.0->google-adk) (2.4.3)\n",
      "Requirement already satisfied: google-resumable-media>=2.7.2 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-cloud-storage<3.0.0,>=2.18.0->google-adk) (2.7.2)\n",
      "Requirement already satisfied: google-crc32c<2.0dev,>=1.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-cloud-storage<3.0.0,>=2.18.0->google-adk) (1.7.1)\n",
      "Requirement already satisfied: certifi in /home/mavrick/anaconda3/lib/python3.12/site-packages (from httpx->opik) (2025.6.15)\n",
      "Requirement already satisfied: httpcore==1.* in /home/mavrick/anaconda3/lib/python3.12/site-packages (from httpx->opik) (1.0.5)\n",
      "Requirement already satisfied: h11<0.15,>=0.13 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from httpcore==1.*->httpx->opik) (0.14.0)\n",
      "Requirement already satisfied: zipp>=3.20 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from importlib-metadata>=6.8.0->litellm) (3.21.0)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from jinja2->opik) (3.0.2)\n",
      "Requirement already satisfied: jsonschema-specifications>=2023.03.6 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from jsonschema<5.0.0,>=4.22.0->litellm) (2023.7.1)\n",
      "Requirement already satisfied: referencing>=0.28.4 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from jsonschema<5.0.0,>=4.22.0->litellm) (0.30.2)\n",
      "Requirement already satisfied: rpds-py>=0.7.1 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from jsonschema<5.0.0,>=4.22.0->litellm) (0.22.3)\n",
      "Requirement already satisfied: httpx-sse>=0.4 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from mcp>=1.8.0->google-adk) (0.4.0)\n",
      "Requirement already satisfied: python-multipart>=0.0.9 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from mcp>=1.8.0->google-adk) (0.0.20)\n",
      "Requirement already satisfied: sse-starlette>=1.6.1 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from mcp>=1.8.0->google-adk) (2.3.6)\n",
      "Requirement already satisfied: distro<2,>=1.7.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from openai<2.0.0->opik) (1.9.0)\n",
      "Requirement already satisfied: jiter<1,>=0.4.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from openai<2.0.0->opik) (0.10.0)\n",
      "Requirement already satisfied: opentelemetry-resourcedetector-gcp==1.*,>=1.5.0dev0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from opentelemetry-exporter-gcp-trace>=1.9.0->google-adk) (1.9.0a0)\n",
      "Requirement already satisfied: opentelemetry-semantic-conventions==0.55b1 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from opentelemetry-sdk>=1.31.0->google-adk) (0.55b1)\n",
      "Requirement already satisfied: annotated-types>=0.6.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from pydantic<3.0.0,>=2.0.0->opik) (0.6.0)\n",
      "Requirement already satisfied: pydantic-core==2.33.2 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from pydantic<3.0.0,>=2.0.0->opik) (2.33.2)\n",
      "Requirement already satisfied: typing-inspection>=0.4.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from pydantic<3.0.0,>=2.0.0->opik) (0.4.1)\n",
      "Requirement already satisfied: six>=1.5 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from python-dateutil>=2.9.0.post0->google-adk) (1.17.0)\n",
      "Requirement already satisfied: charset_normalizer<4,>=2 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from requests>=2.32.4->google-adk) (3.3.2)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from requests>=2.32.4->google-adk) (2.2.3)\n",
      "Requirement already satisfied: greenlet!=0.4.17 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from sqlalchemy>=2.0->google-adk) (3.0.3)\n",
      "Requirement already satisfied: regex>=2022.1.18 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from tiktoken>=0.7.0->litellm) (2024.11.6)\n",
      "Requirement already satisfied: iniconfig in /home/mavrick/anaconda3/lib/python3.12/site-packages (from pytest->opik) (1.1.1)\n",
      "Requirement already satisfied: pluggy<2,>=1.5 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from pytest->opik) (1.5.0)\n",
      "Requirement already satisfied: markdown-it-py>=2.2.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from rich->opik) (2.2.0)\n",
      "Requirement already satisfied: pygments<3.0.0,>=2.13.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from rich->opik) (2.15.1)\n",
      "Requirement already satisfied: huggingface-hub<1.0,>=0.16.4 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from tokenizers->litellm) (0.29.0)\n",
      "Requirement already satisfied: googleapis-common-protos<2.0.0,>=1.56.2 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-api-core!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.0,<3.0.0,>=1.31.5->google-api-python-client>=2.157.0->google-adk) (1.70.0)\n",
      "Requirement already satisfied: grpcio<2.0.0,>=1.33.2 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-api-core[grpc]!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.*,!=2.4.*,!=2.5.*,!=2.6.*,!=2.7.*,<3.0.0,>=1.34.1->google-cloud-aiplatform>=1.95.1->google-cloud-aiplatform[agent-engines]>=1.95.1->google-adk) (1.73.0)\n",
      "Requirement already satisfied: grpcio-status<2.0.0,>=1.33.2 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-api-core[grpc]!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.*,!=2.4.*,!=2.5.*,!=2.6.*,!=2.7.*,<3.0.0,>=1.34.1->google-cloud-aiplatform>=1.95.1->google-cloud-aiplatform[agent-engines]>=1.95.1->google-adk) (1.62.3)\n",
      "Requirement already satisfied: cachetools<6.0,>=2.0.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-auth!=2.24.0,!=2.25.0,<3.0.0,>=1.32.0->google-api-python-client>=2.157.0->google-adk) (5.5.1)\n",
      "Requirement already satisfied: pyasn1-modules>=0.2.1 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-auth!=2.24.0,!=2.25.0,<3.0.0,>=1.32.0->google-api-python-client>=2.157.0->google-adk) (0.2.8)\n",
      "Requirement already satisfied: rsa<5,>=3.1.4 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-auth!=2.24.0,!=2.25.0,<3.0.0,>=1.32.0->google-api-python-client>=2.157.0->google-adk) (4.9.1)\n",
      "Requirement already satisfied: google-cloud-appengine-logging<2.0.0,>=0.1.3 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-cloud-logging<4->google-cloud-aiplatform[agent-engines]>=1.95.1->google-adk) (1.6.2)\n",
      "Requirement already satisfied: google-cloud-audit-log<1.0.0,>=0.3.1 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from google-cloud-logging<4->google-cloud-aiplatform[agent-engines]>=1.95.1->google-adk) (0.3.2)\n",
      "Requirement already satisfied: pyparsing!=3.0.0,!=3.0.1,!=3.0.2,!=3.0.3,<4,>=2.4.2 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from httplib2<1.0.0,>=0.19.0->google-api-python-client>=2.157.0->google-adk) (3.2.0)\n",
      "Requirement already satisfied: filelock in /home/mavrick/anaconda3/lib/python3.12/site-packages (from huggingface-hub<1.0,>=0.16.4->tokenizers->litellm) (3.15.4)\n",
      "Requirement already satisfied: fsspec>=2023.5.0 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from huggingface-hub<1.0,>=0.16.4->tokenizers->litellm) (2024.12.0)\n",
      "Requirement already satisfied: mdurl~=0.1 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from markdown-it-py>=2.2.0->rich->opik) (0.1.0)\n",
      "Requirement already satisfied: numpy<3,>=1.14 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from shapely<3.0.0->google-cloud-aiplatform>=1.95.1->google-cloud-aiplatform[agent-engines]>=1.95.1->google-adk) (1.26.4)\n",
      "Requirement already satisfied: types-awscrt in /home/mavrick/anaconda3/lib/python3.12/site-packages (from botocore-stubs->boto3-stubs>=1.34.110->boto3-stubs[bedrock-runtime]>=1.34.110->opik) (0.27.2)\n",
      "Requirement already satisfied: cffi>=1.12 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from cryptography->authlib>=1.5.1->google-adk) (1.17.1)\n",
      "Requirement already satisfied: pycparser in /home/mavrick/anaconda3/lib/python3.12/site-packages (from cffi>=1.12->cryptography->authlib>=1.5.1->google-adk) (2.21)\n",
      "Requirement already satisfied: pyasn1<0.5.0,>=0.4.6 in /home/mavrick/anaconda3/lib/python3.12/site-packages (from pyasn1-modules>=0.2.1->google-auth!=2.24.0,!=2.25.0,<3.0.0,>=1.32.0->google-api-python-client>=2.157.0->google-adk) (0.4.8)\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install opik google-adk litellm --upgrade"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configure Opik for your session:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "H0DNm-un_0Np"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "OPIK: Opik is already configured. You can check the settings by viewing the config file at /home/mavrick/.opik.config\n"
     ]
    }
   ],
   "source": [
    "import opik\n",
    "opik.configure()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set up your OpenAI API key:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "vN72mHQy_7Ou"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import getpass\n",
    "if \"OPENAI_API_KEY\" not in os.environ:\n",
    "    os.environ[\"OPENAI_API_KEY\"] = getpass.getpass(\"Enter your OpenAI API key: \")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import Required Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import asyncio\n",
    "import datetime\n",
    "from zoneinfo import ZoneInfo\n",
    "\n",
    "from google.adk.agents import LlmAgent\n",
    "from google.adk.models.lite_llm import LiteLlm\n",
    "from google.adk.sessions import InMemorySessionService\n",
    "from google.adk.runners import Runner\n",
    "from google.genai import types\n",
    "from opik.integrations.adk import OpikTracer, track_adk_agent_recursive"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Basic Agent\n",
    "\n",
    "Here we create a single agent with Opik callbacks. The tracer will automatically capture all interactions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_weather(city: str) -> dict:\n",
    "    \"\"\"Get weather information for a city.\"\"\"\n",
    "    if city.lower() == \"new york\":\n",
    "        return {\n",
    "            \"status\": \"success\",\n",
    "            \"report\": \"The weather in New York is sunny with a temperature of 25 °C (77 °F).\",\n",
    "        }\n",
    "    elif city.lower() == \"london\":\n",
    "        return {\n",
    "            \"status\": \"success\",\n",
    "            \"report\": \"The weather in London is cloudy with a temperature of 18 °C (64 °F).\",\n",
    "        }\n",
    "    return {\"status\": \"error\", \"error_message\": f\"Weather info for '{city}' is unavailable.\"}\n",
    "\n",
    "def get_current_time(city: str) -> dict:\n",
    "    \"\"\"Get current time for a city.\"\"\"\n",
    "    if city.lower() == \"new york\":\n",
    "        tz = ZoneInfo(\"America/New_York\")\n",
    "        now = datetime.datetime.now(tz)\n",
    "        return {\n",
    "            \"status\": \"success\",\n",
    "            \"report\": now.strftime(f\"The current time in {city} is %Y-%m-%d %H:%M:%S %Z%z.\"),\n",
    "        }\n",
    "    elif city.lower() == \"london\":\n",
    "        tz = ZoneInfo(\"Europe/London\")\n",
    "        now = datetime.datetime.now(tz)\n",
    "        return {\n",
    "            \"status\": \"success\",\n",
    "            \"report\": now.strftime(f\"The current time in {city} is %Y-%m-%d %H:%M:%S %Z%z.\"),\n",
    "        }\n",
    "    return {\"status\": \"error\", \"error_message\": f\"No timezone info for '{city}'.\"}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configure Opik Tracing\n",
    "\n",
    "Set up the Opik tracer to capture all agent interactions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "basic_tracer = OpikTracer(\n",
    "    name=\"basic-weather-agent\",\n",
    "    tags=[\"basic\", \"weather\", \"time\", \"single-agent\"],\n",
    "    metadata={\n",
    "        \"environment\": \"development\",\n",
    "        \"model\": \"gpt-4o\",\n",
    "        \"framework\": \"google-adk\",\n",
    "        \"example\": \"basic\"\n",
    "    },\n",
    "    project_name=\"adk-basic-demo\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create the LLM Agent\n",
    "\n",
    "Initialize the Google ADK agent with OpenAI's gpt-4o model and Opik tracing:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize LiteLLM with OpenAI gpt-4o\n",
    "llm = LiteLlm(model=\"openai/gpt-4o\")\n",
    "\n",
    "# Create the basic agent with Opik callbacks\n",
    "basic_agent = LlmAgent(\n",
    "    name=\"weather_time_agent\",\n",
    "    model=llm,\n",
    "    description=\"Agent for answering time & weather questions\",\n",
    "    instruction=\"Answer questions about the time or weather in a city. Be helpful and provide clear information.\",\n",
    "    tools=[get_weather, get_current_time],\n",
    "    before_agent_callback=basic_tracer.before_agent_callback,\n",
    "    after_agent_callback=basic_tracer.after_agent_callback,\n",
    "    before_model_callback=basic_tracer.before_model_callback,\n",
    "    after_model_callback=basic_tracer.after_model_callback,\n",
    "    before_tool_callback=basic_tracer.before_tool_callback,\n",
    "    after_tool_callback=basic_tracer.after_tool_callback,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup Session and Runner for Basic Example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "basic_session_service = InMemorySessionService()\n",
    "basic_runner = Runner(\n",
    "    agent=basic_agent,\n",
    "    app_name=\"basic_weather_app\",\n",
    "    session_service=basic_session_service,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Helper Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "async def setup_basic_session():\n",
    "    \"\"\"Create a new session for the basic example.\"\"\"\n",
    "    sess = await basic_session_service.create_session(\n",
    "        app_name=\"basic_weather_app\",\n",
    "        user_id=\"user_basic\",\n",
    "        session_id=\"session_basic_001\"\n",
    "    )\n",
    "    return sess.id\n",
    "\n",
    "async def call_basic_agent(user_msg: str, session_id: str):\n",
    "    \"\"\"Send a message to the basic agent and get the response.\"\"\"\n",
    "    print(f\"User: {user_msg}\")\n",
    "    content = types.Content(role=\"user\", parts=[types.Part(text=user_msg)])\n",
    "    async for event in basic_runner.run_async(user_id=\"user_basic\", session_id=session_id, new_message=content):\n",
    "        if event.is_final_response():\n",
    "            print(f\"Assistant: {event.content.parts[0].text}\")\n",
    "            print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Demo: Basic Agent Interactions\n",
    "\n",
    "Let's test our basic agent with some weather and time queries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Created basic session: session_basic_001\n",
      "\n",
      "User: What's the weather like in New York?\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "OPIK: Started logging traces to the \"adk-basic-demo\" project at https://www.comet.com/opik/api/v1/session/redirect/projects/?trace_id=0197d9b0-c274-72dc-8a1a-fe315109b059&path=aHR0cHM6Ly93d3cuY29tZXQuY29tL29waWsvYXBpLw==.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Assistant: The weather in New York is currently sunny with a temperature of 25°C (77°F).\n",
      "\n",
      "User: What time is it in London?\n",
      "Assistant: The current time in London is 09:25 AM BST.\n",
      "\n",
      "User: Can you tell me both the weather and time in New York?\n",
      "Assistant: In New York, the weather is sunny with a temperature of 25°C (77°F). The current time is 04:25 AM EDT.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Create a session for basic example\n",
    "basic_session_id = await setup_basic_session()\n",
    "print(f\"Created basic session: {basic_session_id}\")\n",
    "print()\n",
    "\n",
    "# Test weather query\n",
    "await call_basic_agent(\"What's the weather like in New York?\", basic_session_id)\n",
    "\n",
    "# Test time query  \n",
    "await call_basic_agent(\"What time is it in London?\", basic_session_id)\n",
    "\n",
    "# Test combined query\n",
    "await call_basic_agent(\"Can you tell me both the weather and time in New York?\", basic_session_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The trace can now be viewed in the UI:\n",
    "\n",
    "![Google Adk Integration Basic Agent](https://raw.githubusercontent.com/comet-ml/opik/main/apps/opik-documentation/documentation/fern/img/cookbook/google_adk_integration_basic_agent.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example 2: Multi-Agent Setup with Hierarchical Tracing\n",
    "\n",
    "This example demonstrates a more complex multi-agent setup where we have specialized agents for different tasks. The key insight is that **you only need to add Opik callbacks to the top-level agent** - all child agent calls will be automatically traced in the same trace tree."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_detailed_weather(city: str) -> dict:\n",
    "    \"\"\"Get detailed weather information including forecast.\"\"\"\n",
    "    weather_data = {\n",
    "        \"new york\": {\n",
    "            \"current\": \"Sunny, 25°C (77°F)\",\n",
    "            \"humidity\": \"65%\",\n",
    "            \"wind\": \"10 km/h NW\",\n",
    "            \"forecast\": \"Partly cloudy tomorrow, high of 27°C\"\n",
    "        },\n",
    "        \"london\": {\n",
    "            \"current\": \"Cloudy, 18°C (64°F)\", \n",
    "            \"humidity\": \"78%\",\n",
    "            \"wind\": \"15 km/h SW\",\n",
    "            \"forecast\": \"Light rain expected tomorrow, high of 16°C\"\n",
    "        },\n",
    "        \"tokyo\": {\n",
    "            \"current\": \"Partly cloudy, 22°C (72°F)\",\n",
    "            \"humidity\": \"70%\", \n",
    "            \"wind\": \"8 km/h E\",\n",
    "            \"forecast\": \"Sunny tomorrow, high of 25°C\"\n",
    "        }\n",
    "    }\n",
    "    \n",
    "    city_lower = city.lower()\n",
    "    if city_lower in weather_data:\n",
    "        data = weather_data[city_lower]\n",
    "        return {\n",
    "            \"status\": \"success\",\n",
    "            \"report\": f\"Weather in {city}: {data['current']}. Humidity: {data['humidity']}, Wind: {data['wind']}. {data['forecast']}\"\n",
    "        }\n",
    "    return {\"status\": \"error\", \"error_message\": f\"Detailed weather for '{city}' is unavailable.\"}\n",
    "\n",
    "def get_world_time(city: str) -> dict:\n",
    "    \"\"\"Get time information for major world cities.\"\"\"\n",
    "    timezones = {\n",
    "        \"new york\": \"America/New_York\",\n",
    "        \"london\": \"Europe/London\", \n",
    "        \"tokyo\": \"Asia/Tokyo\",\n",
    "        \"sydney\": \"Australia/Sydney\",\n",
    "        \"paris\": \"Europe/Paris\"\n",
    "    }\n",
    "    \n",
    "    city_lower = city.lower()\n",
    "    if city_lower in timezones:\n",
    "        tz = ZoneInfo(timezones[city_lower])\n",
    "        now = datetime.datetime.now(tz)\n",
    "        return {\n",
    "            \"status\": \"success\",\n",
    "            \"report\": now.strftime(f\"Current time in {city}: %A, %B %d, %Y at %I:%M %p %Z\")\n",
    "        }\n",
    "    return {\"status\": \"error\", \"error_message\": f\"Time zone info for '{city}' is unavailable.\"}\n",
    "\n",
    "def get_travel_info(from_city: str, to_city: str) -> dict:\n",
    "    \"\"\"Get basic travel information between cities.\"\"\"\n",
    "    travel_data = {\n",
    "        (\"new york\", \"london\"): {\"flight_time\": \"7 hours\", \"time_diff\": \"+5 hours\"},\n",
    "        (\"london\", \"new york\"): {\"flight_time\": \"8 hours\", \"time_diff\": \"-5 hours\"},\n",
    "        (\"new york\", \"tokyo\"): {\"flight_time\": \"14 hours\", \"time_diff\": \"+14 hours\"},\n",
    "        (\"tokyo\", \"new york\"): {\"flight_time\": \"13 hours\", \"time_diff\": \"-14 hours\"},\n",
    "        (\"london\", \"tokyo\"): {\"flight_time\": \"12 hours\", \"time_diff\": \"+9 hours\"},\n",
    "        (\"tokyo\", \"london\"): {\"flight_time\": \"11 hours\", \"time_diff\": \"-9 hours\"},\n",
    "    }\n",
    "    \n",
    "    route = (from_city.lower(), to_city.lower())\n",
    "    if route in travel_data:\n",
    "        data = travel_data[route]\n",
    "        return {\n",
    "            \"status\": \"success\",\n",
    "            \"report\": f\"Travel from {from_city} to {to_city}: Approximately {data['flight_time']} flight time. Time difference: {data['time_diff']}\"\n",
    "        }\n",
    "    return {\"status\": \"error\", \"error_message\": f\"Travel info for '{from_city}' to '{to_city}' is unavailable.\"}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Specialized Agents\n",
    "\n",
    "Now we'll create specialized agents for different domains. Notice that only the coordinator agent will have Opik callbacks:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Weather specialist agent (no Opik callbacks needed)\n",
    "weather_agent = LlmAgent(\n",
    "    name=\"weather_specialist\",\n",
    "    model=llm,\n",
    "    description=\"Specialized agent for detailed weather information\",\n",
    "    instruction=\"Provide comprehensive weather information including current conditions and forecasts. Be detailed and informative.\",\n",
    "    tools=[get_detailed_weather]\n",
    ")\n",
    "\n",
    "# Time specialist agent (no Opik callbacks needed)  \n",
    "time_agent = LlmAgent(\n",
    "    name=\"time_specialist\",\n",
    "    model=llm,\n",
    "    description=\"Specialized agent for world time information\",\n",
    "    instruction=\"Provide accurate time information for cities around the world. Include day of week and full date.\",\n",
    "    tools=[get_world_time]\n",
    ")\n",
    "\n",
    "# Travel specialist agent (no Opik callbacks needed)\n",
    "travel_agent = LlmAgent(\n",
    "    name=\"travel_specialist\", \n",
    "    model=llm,\n",
    "    description=\"Specialized agent for travel information\",\n",
    "    instruction=\"Provide helpful travel information including flight times and time zone differences.\",\n",
    "    tools=[get_travel_info]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Coordinator Agent with Opik Tracing\n",
    "\n",
    "The coordinator agent orchestrates the specialized agents and **only needs Opik callbacks here** - all child agent calls will be automatically traced:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "OPIK: `track_adk_agent_recursive` is experimental feature. Please let us know if something is not working as expected: https://github.com/comet-ml/opik/issues\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "LlmAgent(name='travel_coordinator', description='Coordinator agent that delegates to specialized agents for weather, time, and travel information', parent_agent=None, sub_agents=[LlmAgent(name='weather_specialist', description='Specialized agent for detailed weather information', parent_agent=LlmAgent(name='travel_coordinator', description='Coordinator agent that delegates to specialized agents for weather, time, and travel information', parent_agent=None, sub_agents=[...], before_agent_callback=<bound method OpikTracer.before_agent_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>, after_agent_callback=<bound method OpikTracer.after_agent_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>, model=LiteLlm(model='openai/gpt-4o', llm_client=<google.adk.models.lite_llm.LiteLLMClient object at 0x7f2ef2041df0>), instruction='You are a travel coordinator that helps users with weather, time, and travel information.\\n    \\n    You have access to three specialized agents:\\n    - weather_specialist: For detailed weather information\\n    - time_specialist: For world time information  \\n    - travel_specialist: For travel planning information\\n    \\n    Delegate appropriate queries to the right specialist agents and compile comprehensive responses for the user.', global_instruction='', tools=[], generate_content_config=None, disallow_transfer_to_parent=False, disallow_transfer_to_peers=False, include_contents='default', input_schema=None, output_schema=None, output_key=None, planner=None, code_executor=None, before_model_callback=<bound method OpikTracer.before_model_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>, after_model_callback=<bound method OpikTracer.after_model_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>, before_tool_callback=<bound method OpikTracer.before_tool_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>, after_tool_callback=<bound method OpikTracer.after_tool_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>), sub_agents=[], before_agent_callback=<bound method OpikTracer.before_agent_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>, after_agent_callback=<bound method OpikTracer.after_agent_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>, model=LiteLlm(model='openai/gpt-4o', llm_client=<google.adk.models.lite_llm.LiteLLMClient object at 0x7f2ef2041df0>), instruction='Provide comprehensive weather information including current conditions and forecasts. Be detailed and informative.', global_instruction='', tools=[<function get_detailed_weather at 0x7f2ef6070a40>], generate_content_config=None, disallow_transfer_to_parent=False, disallow_transfer_to_peers=False, include_contents='default', input_schema=None, output_schema=None, output_key=None, planner=None, code_executor=None, before_model_callback=<bound method OpikTracer.before_model_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>, after_model_callback=<bound method OpikTracer.after_model_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>, before_tool_callback=<bound method OpikTracer.before_tool_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>, after_tool_callback=<bound method OpikTracer.after_tool_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>), LlmAgent(name='time_specialist', description='Specialized agent for world time information', parent_agent=LlmAgent(name='travel_coordinator', description='Coordinator agent that delegates to specialized agents for weather, time, and travel information', parent_agent=None, sub_agents=[...], before_agent_callback=<bound method OpikTracer.before_agent_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>, after_agent_callback=<bound method OpikTracer.after_agent_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>, model=LiteLlm(model='openai/gpt-4o', llm_client=<google.adk.models.lite_llm.LiteLLMClient object at 0x7f2ef2041df0>), instruction='You are a travel coordinator that helps users with weather, time, and travel information.\\n    \\n    You have access to three specialized agents:\\n    - weather_specialist: For detailed weather information\\n    - time_specialist: For world time information  \\n    - travel_specialist: For travel planning information\\n    \\n    Delegate appropriate queries to the right specialist agents and compile comprehensive responses for the user.', global_instruction='', tools=[], generate_content_config=None, disallow_transfer_to_parent=False, disallow_transfer_to_peers=False, include_contents='default', input_schema=None, output_schema=None, output_key=None, planner=None, code_executor=None, before_model_callback=<bound method OpikTracer.before_model_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>, after_model_callback=<bound method OpikTracer.after_model_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>, before_tool_callback=<bound method OpikTracer.before_tool_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>, after_tool_callback=<bound method OpikTracer.after_tool_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>), sub_agents=[], before_agent_callback=<bound method OpikTracer.before_agent_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>, after_agent_callback=<bound method OpikTracer.after_agent_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>, model=LiteLlm(model='openai/gpt-4o', llm_client=<google.adk.models.lite_llm.LiteLLMClient object at 0x7f2ef2041df0>), instruction='Provide accurate time information for cities around the world. Include day of week and full date.', global_instruction='', tools=[<function get_world_time at 0x7f2ef60707c0>], generate_content_config=None, disallow_transfer_to_parent=False, disallow_transfer_to_peers=False, include_contents='default', input_schema=None, output_schema=None, output_key=None, planner=None, code_executor=None, before_model_callback=<bound method OpikTracer.before_model_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>, after_model_callback=<bound method OpikTracer.after_model_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>, before_tool_callback=<bound method OpikTracer.before_tool_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>, after_tool_callback=<bound method OpikTracer.after_tool_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>), LlmAgent(name='travel_specialist', description='Specialized agent for travel information', parent_agent=LlmAgent(name='travel_coordinator', description='Coordinator agent that delegates to specialized agents for weather, time, and travel information', parent_agent=None, sub_agents=[...], before_agent_callback=<bound method OpikTracer.before_agent_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>, after_agent_callback=<bound method OpikTracer.after_agent_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>, model=LiteLlm(model='openai/gpt-4o', llm_client=<google.adk.models.lite_llm.LiteLLMClient object at 0x7f2ef2041df0>), instruction='You are a travel coordinator that helps users with weather, time, and travel information.\\n    \\n    You have access to three specialized agents:\\n    - weather_specialist: For detailed weather information\\n    - time_specialist: For world time information  \\n    - travel_specialist: For travel planning information\\n    \\n    Delegate appropriate queries to the right specialist agents and compile comprehensive responses for the user.', global_instruction='', tools=[], generate_content_config=None, disallow_transfer_to_parent=False, disallow_transfer_to_peers=False, include_contents='default', input_schema=None, output_schema=None, output_key=None, planner=None, code_executor=None, before_model_callback=<bound method OpikTracer.before_model_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>, after_model_callback=<bound method OpikTracer.after_model_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>, before_tool_callback=<bound method OpikTracer.before_tool_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>, after_tool_callback=<bound method OpikTracer.after_tool_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>), sub_agents=[], before_agent_callback=<bound method OpikTracer.before_agent_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>, after_agent_callback=<bound method OpikTracer.after_agent_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>, model=LiteLlm(model='openai/gpt-4o', llm_client=<google.adk.models.lite_llm.LiteLLMClient object at 0x7f2ef2041df0>), instruction='Provide helpful travel information including flight times and time zone differences.', global_instruction='', tools=[<function get_travel_info at 0x7f2ef6070720>], generate_content_config=None, disallow_transfer_to_parent=False, disallow_transfer_to_peers=False, include_contents='default', input_schema=None, output_schema=None, output_key=None, planner=None, code_executor=None, before_model_callback=<bound method OpikTracer.before_model_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>, after_model_callback=<bound method OpikTracer.after_model_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>, before_tool_callback=<bound method OpikTracer.before_tool_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>, after_tool_callback=<bound method OpikTracer.after_tool_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>)], before_agent_callback=<bound method OpikTracer.before_agent_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>, after_agent_callback=<bound method OpikTracer.after_agent_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>, model=LiteLlm(model='openai/gpt-4o', llm_client=<google.adk.models.lite_llm.LiteLLMClient object at 0x7f2ef2041df0>), instruction='You are a travel coordinator that helps users with weather, time, and travel information.\\n    \\n    You have access to three specialized agents:\\n    - weather_specialist: For detailed weather information\\n    - time_specialist: For world time information  \\n    - travel_specialist: For travel planning information\\n    \\n    Delegate appropriate queries to the right specialist agents and compile comprehensive responses for the user.', global_instruction='', tools=[], generate_content_config=None, disallow_transfer_to_parent=False, disallow_transfer_to_peers=False, include_contents='default', input_schema=None, output_schema=None, output_key=None, planner=None, code_executor=None, before_model_callback=<bound method OpikTracer.before_model_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>, after_model_callback=<bound method OpikTracer.after_model_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>, before_tool_callback=<bound method OpikTracer.before_tool_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>, after_tool_callback=<bound method OpikTracer.after_tool_callback of <opik.integrations.adk.opik_tracer.OpikTracer object at 0x7f2ef04364e0>>)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Configure Opik tracer for multi-agent example\n",
    "multi_agent_tracer = OpikTracer(\n",
    "    name=\"multi-agent-coordinator\",\n",
    "    tags=[\"multi-agent\", \"coordinator\", \"weather\", \"time\", \"travel\"],\n",
    "    metadata={\n",
    "        \"environment\": \"development\",\n",
    "        \"model\": \"gpt-4o\",\n",
    "        \"framework\": \"google-adk\", \n",
    "        \"example\": \"multi-agent\",\n",
    "        \"agent_count\": 4\n",
    "    },\n",
    "    project_name=\"adk-multi-agent-demo\"\n",
    ")\n",
    "\n",
    "# Coordinator agent with sub-agents\n",
    "coordinator_agent = LlmAgent(\n",
    "    name=\"travel_coordinator\",\n",
    "    model=llm,\n",
    "    description=\"Coordinator agent that delegates to specialized agents for weather, time, and travel information\",\n",
    "    instruction=\"\"\"You are a travel coordinator that helps users with weather, time, and travel information.\n",
    "    \n",
    "    You have access to three specialized agents:\n",
    "    - weather_specialist: For detailed weather information\n",
    "    - time_specialist: For world time information  \n",
    "    - travel_specialist: For travel planning information\n",
    "    \n",
    "    Delegate appropriate queries to the right specialist agents and compile comprehensive responses for the user.\"\"\",\n",
    "    tools=[],  # No direct tools, delegates to sub-agents\n",
    "    sub_agents=[weather_agent, time_agent, travel_agent],\n",
    ")\n",
    "\n",
    "# Use the experimental recursive tracking feature to instrument all agents at once\n",
    "from opik.integrations.adk import track_adk_agent_recursive\n",
    "track_adk_agent_recursive(coordinator_agent, multi_agent_tracer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup Multi-Agent Session and Runner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "multi_session_service = InMemorySessionService()\n",
    "multi_runner = Runner(\n",
    "    agent=coordinator_agent,\n",
    "    app_name=\"multi_agent_travel_app\",\n",
    "    session_service=multi_session_service,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Multi-Agent Helper Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "async def setup_multi_session():\n",
    "    \"\"\"Create a new session for the multi-agent example.\"\"\"\n",
    "    sess = await multi_session_service.create_session(\n",
    "        app_name=\"multi_agent_travel_app\",\n",
    "        user_id=\"user_multi\",\n",
    "        session_id=\"session_multi_001\"\n",
    "    )\n",
    "    return sess.id\n",
    "\n",
    "async def call_multi_agent(user_msg: str, session_id: str):\n",
    "    \"\"\"Send a message to the coordinator agent and get the response.\"\"\"\n",
    "    print(f\"User: {user_msg}\")\n",
    "    content = types.Content(role=\"user\", parts=[types.Part(text=user_msg)])\n",
    "    async for event in multi_runner.run_async(user_id=\"user_multi\", session_id=session_id, new_message=content):\n",
    "        if event.is_final_response():\n",
    "            print(f\"Coordinator: {event.content.parts[0].text}\")\n",
    "            print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Demo: Multi-Agent Interactions\n",
    "\n",
    "Let's test our multi-agent setup with complex queries that require coordination between agents:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=== Creating separate sessions for individual traces ===\n",
      "User: I need detailed weather information for Tokyo\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "OPIK: Started logging traces to the \"adk-multi-agent-demo\" project at https://www.comet.com/opik/api/v1/session/redirect/projects/?trace_id=0197d9b0-f290-791a-82f4-7b3673b6c72f&path=aHR0cHM6Ly93d3cuY29tZXQuY29tL29waWsvYXBpLw==.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Coordinator: The current weather in Tokyo is partly cloudy with a temperature of 22°C (72°F). The humidity is at 70%, and there's an east wind blowing at 8 km/h. The forecast for tomorrow predicts sunny weather with a high of 25°C.\n",
      "\n",
      "User: What time is it in Paris right now?\n",
      "Coordinator: The current time in Paris is Saturday, July 5, 2025, at 10:25 AM CEST.\n",
      "\n",
      "User: I'm planning to travel from London to New York. Can you help with travel time and time zones?\n",
      "Coordinator: The current time information is as follows:\n",
      "\n",
      "- **London, UK**: Saturday, July 5, 2025, at 09:25 AM BST (British Summer Time)\n",
      "- **New York, USA**: Saturday, July 5, 2025, at 04:25 AM EDT (Eastern Daylight Time)\n",
      "\n",
      "Please note that London is currently 5 hours ahead of New York. If you need additional travel information, I recommend reaching out to the travel_specialist agent.\n",
      "\n",
      "User: I'm traveling from New York to Tokyo tomorrow. Can you give me the weather in both cities, current times, and travel information?\n",
      "Coordinator: Here is the current time information for your trip from New York to Tokyo:\n",
      "\n",
      "- **New York**: Saturday, July 05, 2025, at 04:25 AM EDT\n",
      "- **Tokyo**: Saturday, July 05, 2025, at 05:25 PM JST\n",
      "\n",
      "The travel time is approximately 14 hours, with a time difference of +14 hours from New York to Tokyo.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Demo: Each question in a separate session to show individual traces\n",
    "print(\"=== Creating separate sessions for individual traces ===\")\n",
    "\n",
    "# Question 1: Weather (separate session)\n",
    "session_1 = await setup_multi_session()\n",
    "await call_multi_agent(\"I need detailed weather information for Tokyo\", session_1)\n",
    "\n",
    "# Question 2: Time (separate session) \n",
    "session_2 = await setup_multi_session()\n",
    "await call_multi_agent(\"What time is it in Paris right now?\", session_2)\n",
    "\n",
    "# Question 3: Travel (separate session)\n",
    "session_3 = await setup_multi_session()\n",
    "await call_multi_agent(\"I'm planning to travel from London to New York. Can you help with travel time and time zones?\", session_3)\n",
    "\n",
    "# Question 4: Complex multi-agent (separate session)\n",
    "session_4 = await setup_multi_session()\n",
    "await call_multi_agent(\"I'm traveling from New York to Tokyo tomorrow. Can you give me the weather in both cities, current times, and travel information?\", session_4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The trace can now be viewed in the UI:\n",
    "\n",
    "![Google Adk Integration Multi Agent](https://raw.githubusercontent.com/comet-ml/opik/main/apps/opik-documentation/documentation/fern/img/cookbook/google_adk_integration_multi_agent.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example 3: Hybrid Tracing - Combining Opik Decorators with ADK Callbacks\n",
    "\n",
    "This advanced example shows how to combine Opik's `@opik.track` decorator with ADK's callback system. This is powerful when you have complex multi-step tools that perform their own internal operations that you want to trace separately, while still maintaining the overall agent trace context.\n",
    "\n",
    "### Define Advanced Tools with Opik Decorators\n",
    "\n",
    "These tools use the `@opik.track` decorator to trace their internal operations, while still being called within the ADK agent's trace context:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "@opik.track(name=\"weather_data_processing\", tags=[\"data-processing\", \"weather\"])\n",
    "def process_weather_data(raw_data: dict) -> dict:\n",
    "    \"\"\"Process raw weather data with additional computations.\"\"\"\n",
    "    # Simulate some data processing steps that we want to trace separately\n",
    "    processed = {\n",
    "        \"temperature_celsius\": raw_data.get(\"temp_c\", 0),\n",
    "        \"temperature_fahrenheit\": raw_data.get(\"temp_c\", 0) * 9/5 + 32,\n",
    "        \"conditions\": raw_data.get(\"condition\", \"unknown\"),\n",
    "        \"comfort_index\": \"comfortable\" if 18 <= raw_data.get(\"temp_c\", 0) <= 25 else \"less comfortable\"\n",
    "    }\n",
    "    return processed\n",
    "\n",
    "@opik.track(name=\"location_validation\", tags=[\"validation\", \"location\"])  \n",
    "def validate_location(city: str) -> dict:\n",
    "    \"\"\"Validate and normalize city names.\"\"\"\n",
    "    # Simulate location validation logic that we want to trace\n",
    "    normalized_cities = {\n",
    "        \"nyc\": \"New York\",\n",
    "        \"ny\": \"New York\", \n",
    "        \"new york city\": \"New York\",\n",
    "        \"london uk\": \"London\",\n",
    "        \"london england\": \"London\",\n",
    "        \"tokyo japan\": \"Tokyo\"\n",
    "    }\n",
    "    \n",
    "    city_lower = city.lower().strip()\n",
    "    validated_city = normalized_cities.get(city_lower, city.title())\n",
    "    \n",
    "    return {\n",
    "        \"original\": city,\n",
    "        \"validated\": validated_city,\n",
    "        \"is_valid\": city_lower in [\"new york\", \"london\", \"tokyo\"] or city_lower in normalized_cities\n",
    "    }\n",
    "\n",
    "@opik.track(name=\"advanced_weather_lookup\", tags=[\"weather\", \"api-simulation\"])\n",
    "def get_advanced_weather(city: str) -> dict:\n",
    "    \"\"\"Get weather with internal processing steps tracked by Opik decorators.\"\"\"\n",
    "    \n",
    "    # Step 1: Validate location (traced by @opik.track)\n",
    "    location_result = validate_location(city)\n",
    "    \n",
    "    if not location_result[\"is_valid\"]:\n",
    "        return {\n",
    "            \"status\": \"error\", \n",
    "            \"error_message\": f\"Invalid location: {city}\"\n",
    "        }\n",
    "    \n",
    "    validated_city = location_result[\"validated\"]\n",
    "    \n",
    "    # Step 2: Get raw weather data (simulated)\n",
    "    raw_weather_data = {\n",
    "        \"New York\": {\"temp_c\": 25, \"condition\": \"sunny\", \"humidity\": 65},\n",
    "        \"London\": {\"temp_c\": 18, \"condition\": \"cloudy\", \"humidity\": 78}, \n",
    "        \"Tokyo\": {\"temp_c\": 22, \"condition\": \"partly cloudy\", \"humidity\": 70}\n",
    "    }\n",
    "    \n",
    "    if validated_city not in raw_weather_data:\n",
    "        return {\n",
    "            \"status\": \"error\",\n",
    "            \"error_message\": f\"Weather data unavailable for {validated_city}\"\n",
    "        }\n",
    "    \n",
    "    raw_data = raw_weather_data[validated_city]\n",
    "    \n",
    "    # Step 3: Process the data (traced by @opik.track)\n",
    "    processed_data = process_weather_data(raw_data)\n",
    "    \n",
    "    return {\n",
    "        \"status\": \"success\",\n",
    "        \"city\": validated_city,\n",
    "        \"report\": f\"Weather in {validated_city}: {processed_data['conditions']}, {processed_data['temperature_celsius']}°C ({processed_data['temperature_fahrenheit']:.1f}°F). Comfort level: {processed_data['comfort_index']}.\",\n",
    "        \"raw_humidity\": raw_data[\"humidity\"]\n",
    "    }\n",
    "\n",
    "@opik.track(name=\"time_calculation\", tags=[\"time\", \"calculation\"])\n",
    "def calculate_time_info(timezone_name: str) -> dict:\n",
    "    \"\"\"Calculate detailed time information with internal steps.\"\"\"\n",
    "    try:\n",
    "        tz = ZoneInfo(timezone_name)\n",
    "        now = datetime.datetime.now(tz)\n",
    "        \n",
    "        # Calculate additional time info\n",
    "        time_info = {\n",
    "            \"current_time\": now,\n",
    "            \"hour_24\": now.hour,\n",
    "            \"is_business_hours\": 9 <= now.hour <= 17,\n",
    "            \"day_of_week\": now.strftime(\"%A\"),\n",
    "            \"week_number\": now.isocalendar()[1],\n",
    "            \"is_weekend\": now.weekday() >= 5\n",
    "        }\n",
    "        \n",
    "        return time_info\n",
    "    except Exception as e:\n",
    "        return {\"error\": str(e)}\n",
    "\n",
    "@opik.track(name=\"advanced_time_lookup\", tags=[\"time\", \"timezone\"])  \n",
    "def get_advanced_time(city: str) -> dict:\n",
    "    \"\"\"Get time information with advanced calculations.\"\"\"\n",
    "    \n",
    "    # Step 1: Validate location\n",
    "    location_result = validate_location(city)\n",
    "    \n",
    "    if not location_result[\"is_valid\"]:\n",
    "        return {\n",
    "            \"status\": \"error\",\n",
    "            \"error_message\": f\"Invalid location for time lookup: {city}\"\n",
    "        }\n",
    "    \n",
    "    validated_city = location_result[\"validated\"]\n",
    "    \n",
    "    # Step 2: Map to timezone\n",
    "    timezone_map = {\n",
    "        \"New York\": \"America/New_York\",\n",
    "        \"London\": \"Europe/London\",\n",
    "        \"Tokyo\": \"Asia/Tokyo\"\n",
    "    }\n",
    "    \n",
    "    if validated_city not in timezone_map:\n",
    "        return {\n",
    "            \"status\": \"error\", \n",
    "            \"error_message\": f\"Timezone mapping unavailable for {validated_city}\"\n",
    "        }\n",
    "    \n",
    "    # Step 3: Calculate time info (traced by @opik.track)\n",
    "    time_info = calculate_time_info(timezone_map[validated_city])\n",
    "    \n",
    "    if \"error\" in time_info:\n",
    "        return {\"status\": \"error\", \"error_message\": time_info[\"error\"]}\n",
    "    \n",
    "    current_time = time_info[\"current_time\"]\n",
    "    business_status = \"during business hours\" if time_info[\"is_business_hours\"] else \"outside business hours\"\n",
    "    weekend_status = \"on a weekend\" if time_info[\"is_weekend\"] else \"on a weekday\"\n",
    "    \n",
    "    return {\n",
    "        \"status\": \"success\",\n",
    "        \"city\": validated_city,\n",
    "        \"report\": f\"Current time in {validated_city}: {current_time.strftime('%A, %B %d, %Y at %I:%M %p %Z')} ({business_status}, {weekend_status})\",\n",
    "        \"metadata\": {\n",
    "            \"hour_24\": time_info[\"hour_24\"],\n",
    "            \"week_number\": time_info[\"week_number\"],\n",
    "            \"is_business_hours\": time_info[\"is_business_hours\"]\n",
    "        }\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Hybrid Agent with Both Decorator and Callback Tracing\n",
    "\n",
    "This agent uses tools that have internal Opik tracing via decorators, while the agent itself uses ADK callbacks:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configure Opik tracer for hybrid example\n",
    "hybrid_tracer = OpikTracer(\n",
    "    name=\"hybrid-tracing-agent\",\n",
    "    tags=[\"hybrid\", \"decorators\", \"callbacks\", \"advanced\"],\n",
    "    metadata={\n",
    "        \"environment\": \"development\",\n",
    "        \"model\": \"gpt-4o\",\n",
    "        \"framework\": \"google-adk\",\n",
    "        \"example\": \"hybrid-tracing\",\n",
    "        \"tracing_methods\": [\"decorators\", \"callbacks\"]\n",
    "    },\n",
    "    project_name=\"adk-hybrid-demo\"\n",
    ")\n",
    "\n",
    "# Create hybrid agent that combines both tracing approaches\n",
    "hybrid_agent = LlmAgent(\n",
    "    name=\"advanced_weather_time_agent\",\n",
    "    model=llm,\n",
    "    description=\"Advanced agent with hybrid Opik tracing using both decorators and callbacks\",\n",
    "    instruction=\"\"\"You are an advanced weather and time agent that provides detailed information with comprehensive internal processing.\n",
    "    \n",
    "    Your tools perform multi-step operations that are individually traced, giving detailed visibility into the processing pipeline.\n",
    "    Use the advanced weather and time tools to provide thorough, well-processed information to users.\"\"\",\n",
    "    tools=[get_advanced_weather, get_advanced_time],\n",
    "    # ADK callbacks for agent-level tracing\n",
    "    before_agent_callback=hybrid_tracer.before_agent_callback,\n",
    "    after_agent_callback=hybrid_tracer.after_agent_callback,\n",
    "    before_model_callback=hybrid_tracer.before_model_callback,\n",
    "    after_model_callback=hybrid_tracer.after_model_callback,\n",
    "    before_tool_callback=hybrid_tracer.before_tool_callback,\n",
    "    after_tool_callback=hybrid_tracer.after_tool_callback,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup Hybrid Session and Runner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "hybrid_session_service = InMemorySessionService()\n",
    "hybrid_runner = Runner(\n",
    "    agent=hybrid_agent,\n",
    "    app_name=\"hybrid_tracing_app\",\n",
    "    session_service=hybrid_session_service,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hybrid Example Helper Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "async def setup_hybrid_session():\n",
    "    \"\"\"Create a new session for the hybrid tracing example.\"\"\"\n",
    "    sess = await hybrid_session_service.create_session(\n",
    "        app_name=\"hybrid_tracing_app\",\n",
    "        user_id=\"user_hybrid\",\n",
    "        session_id=\"session_hybrid_001\"\n",
    "    )\n",
    "    return sess.id\n",
    "\n",
    "async def call_hybrid_agent(user_msg: str, session_id: str):\n",
    "    \"\"\"Send a message to the hybrid agent and get the response.\"\"\"\n",
    "    print(f\"User: {user_msg}\")\n",
    "    content = types.Content(role=\"user\", parts=[types.Part(text=user_msg)])\n",
    "    async for event in hybrid_runner.run_async(user_id=\"user_hybrid\", session_id=session_id, new_message=content):\n",
    "        if event.is_final_response():\n",
    "            print(f\"Advanced Assistant: {event.content.parts[0].text}\")\n",
    "            print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "### Demo: Hybrid Tracing in Action\n",
    "\n",
    "This demo showcases how both decorator-traced tool operations and ADK callback-traced agent operations appear in the same unified trace:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Created hybrid tracing session: session_hybrid_001\n",
      "\n",
      "=== Testing Advanced Weather Lookup ===\n",
      "User: What's the weather like in NYC?\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "OPIK: Started logging traces to the \"adk-hybrid-demo\" project at https://www.comet.com/opik/api/v1/session/redirect/projects/?trace_id=0197d9b1-3310-7592-b309-ba8143a52037&path=aHR0cHM6Ly93d3cuY29tZXQuY29tL29waWsvYXBpLw==.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Advanced Assistant: The weather in New York City is currently sunny with a temperature of 25°C (77.0°F). The comfort level is described as comfortable, with a humidity level of approximately 65%.\n",
      "\n",
      "=== Testing Advanced Time Lookup ===\n",
      "User: What time is it in London right now? Include business hours info.\n",
      "Advanced Assistant: The current time in London is 09:26 AM BST on Saturday, July 5, 2025. Although it is during business hours, it's worth noting that today is a weekend.\n",
      "\n",
      "=== Testing Location Validation ===\n",
      "User: What's the weather in InvalidCity?\n",
      "Advanced Assistant: I couldn't retrieve weather information for \"InvalidCity\" as it seems to be an invalid or unrecognized location. Please check the city name and try again.\n",
      "\n",
      "=== Testing Complex Multi-Step Query ===\n",
      "User: I need detailed weather and time information for Tokyo, including whether it's business hours\n",
      "Advanced Assistant: Here's the detailed information for Tokyo:\n",
      "\n",
      "- **Weather**: The weather in Tokyo is partly cloudy with a temperature of 22°C (71.6°F). The comfort level is described as comfortable, with a humidity level of approximately 70%.\n",
      "\n",
      "- **Time**: The current time in Tokyo is 05:26 PM JST on Saturday, July 5, 2025. Although it's during business hours, it is important to note that today is the weekend.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Create a session for hybrid example\n",
    "hybrid_session_id = await setup_hybrid_session()\n",
    "print(f\"Created hybrid tracing session: {hybrid_session_id}\")\n",
    "print()\n",
    "\n",
    "# Test weather with internal processing steps\n",
    "print(\"=== Testing Advanced Weather Lookup ===\")\n",
    "await call_hybrid_agent(\"What's the weather like in NYC?\", hybrid_session_id)\n",
    "\n",
    "# Test time with internal calculations\n",
    "print(\"=== Testing Advanced Time Lookup ===\") \n",
    "await call_hybrid_agent(\"What time is it in London right now? Include business hours info.\", hybrid_session_id)\n",
    "\n",
    "# Test invalid location to see validation tracing\n",
    "print(\"=== Testing Location Validation ===\")\n",
    "await call_hybrid_agent(\"What's the weather in InvalidCity?\", hybrid_session_id)\n",
    "\n",
    "# Test complex query that triggers multiple internal operations\n",
    "print(\"=== Testing Complex Multi-Step Query ===\")\n",
    "await call_hybrid_agent(\"I need detailed weather and time information for Tokyo, including whether it's business hours\", hybrid_session_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The trace can now be viewed in the UI:\n",
    "\n",
    "![Google Adk Integration Hybrid Agent](https://raw.githubusercontent.com/comet-ml/opik/main/apps/opik-documentation/documentation/fern/img/cookbook/google_adk_integration_hybrid_agent.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Understanding the Tracing Output\n",
    "\n",
    "After running all three examples, you can view the traces in your Opik dashboard. Here's what you'll see:\n",
    "\n",
    "### Basic Example Traces\n",
    "- Simple linear trace showing: User Input → Agent Processing → Tool Calls → Model Response\n",
    "- Clear visibility into tool execution and model interactions\n",
    "\n",
    "### Multi-Agent Example Traces  \n",
    "- Hierarchical trace showing: Coordinator Agent → Sub-Agent Delegation → Specialized Tool Calls\n",
    "- **Key insight**: Only the coordinator needed Opik callbacks, but all sub-agent operations are traced\n",
    "- Shows the complete decision tree of which agents were involved\n",
    "\n",
    "### Hybrid Example Traces\n",
    "- **Most comprehensive**: Shows both ADK callback traces AND decorator traces in the same trace tree\n",
    "- Tool calls contain nested spans from `@opik.track` decorators\n",
    "- Demonstrates how decorator-traced functions (like `validate_location`, `process_weather_data`) appear as child spans within the tool call spans\n",
    "- Perfect for debugging complex multi-step operations\n",
    "\n",
    "## Key Benefits of This Integration\n",
    "\n",
    "1. **Automatic Tracing**: ADK callbacks provide zero-configuration tracing of agent interactions\n",
    "2. **Hierarchical Visibility**: Multi-agent setups automatically create nested trace structures  \n",
    "3. **Flexible Granularity**: Combine coarse-grained agent tracing with fine-grained function tracing\n",
    "4. **Unified Context**: All traces (callbacks + decorators) appear in the same trace tree\n",
    "5. **Production Ready**: Comprehensive observability for debugging and optimization\n",
    "\n",
    "## Next Steps\n",
    "\n",
    "This notebook demonstrated three powerful integration patterns. You can extend this by:\n",
    "\n",
    "- Adding custom evaluation metrics using Opik's evaluation framework\n",
    "- Implementing real-time monitoring and alerting based on trace data\n",
    "- Using different LLM models and comparing their performance\n",
    "- Adding more sophisticated multi-agent workflows\n",
    "- Implementing custom tracing strategies for specific business logic\n",
    "- Building evaluation datasets from traced conversations\n",
    "\n",
    "For more information:\n",
    "- [Google ADK Documentation](https://google.github.io/adk-docs/)\n",
    "- [Opik Documentation](https://www.comet.com/docs/opik/)\n",
    "- [LiteLLM Documentation](https://docs.litellm.ai/)\n",
    "- [Opik Integrations Guide](https://www.comet.com/docs/opik/integrations/)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
